#include <iostream>
#include <array>

#ifdef WAFER_OPENSSL
#include <openssl/evp.h>
#endif

#include "container.h"

#define MSG_LEN 16
#define MSG_CNT 100

#define AES_DIGEST_LEN MSG_LEN
#define INPUT_SIZE (MSG_LEN * MSG_CNT)
#define OUTPUT_SIZE (AES_DIGEST_LEN * MSG_CNT)
#define KEY_SIZE 32

enum class AES_TYPE {
  AES128 = 16,
  AES192 = 24,
  AES256 = 32
};

#ifndef WAFER_OPENSSL
extern "C"  {
    void CUDA(MemRef<unsigned char, 1>* inputs, int msgLen, MemRef<unsigned char, 1>* outputs, int cptLen, MemRef<unsigned char, 1>* key, int keyLen, int msgCnt);
}

void AES_CRACK_OP(MemRef<unsigned char, 1>* inputs, int msgLen, MemRef<unsigned char, 1>* outputs, int cptLen, MemRef<unsigned char, 1>* key, int keyLen, int msgCnt) {
    CUDA(inputs, msgLen, outputs, cptLen, key, keyLen, msgCnt);
}
#endif

uint8_t exampleInput[INPUT_SIZE] = {0xE2, 0x4C, 0xF1, 0xEC, 0x3C, 0x18, 0xAD, 0x6E, 0x8E, 0xB8, 0x45, 0xEA, 0xD1, 0x34, 0xAF, 0xE2, 
                                    0xD5, 0xAB, 0xCA, 0x17, 0xC2, 0x45, 0x04, 0x02, 0x76, 0x28, 0xC5, 0x57, 0x05, 0xC6, 0x46, 0x38, 
                                    0x45, 0xD2, 0x00, 0xE3, 0xB1, 0xD9, 0x28, 0x6B, 0x92, 0x7D, 0xEC, 0xA0, 0xB6, 0x56, 0xFA, 0x92, 
                                    0x88, 0xB5, 0x96, 0x2B, 0xB3, 0x60, 0xF4, 0x2A, 0x83, 0x67, 0x6C, 0x21, 0x5B, 0xC8, 0x61, 0x5B, 
                                    0x96, 0x7E, 0x12, 0x39, 0x06, 0xC8, 0xB5, 0xD3, 0x0F, 0xC2, 0x25, 0xB9, 0xD1, 0x39, 0x67, 0x5A, 
                                    0x69, 0x9D, 0xD6, 0x07, 0x0B, 0xBF, 0x66, 0xD6, 0xAF, 0xD1, 0x94, 0x42, 0x2F, 0x47, 0x9F, 0x04, 
                                    0xCA, 0x85, 0xC7, 0x28, 0x36, 0xC1, 0x86, 0x48, 0xF4, 0x05, 0xD9, 0xBE, 0x44, 0x3F, 0x30, 0x8C, 
                                    0xFA, 0x2E, 0xB6, 0x4E, 0x65, 0x56, 0x15, 0x58, 0x56, 0x62, 0x19, 0x5D, 0xD7, 0xFD, 0xE9, 0x66, 
                                    0x59, 0x0E, 0x4B, 0xDD, 0xE9, 0x39, 0x3D, 0x10, 0x55, 0xE1, 0xA5, 0xB9, 0x0B, 0x42, 0x0B, 0x39, 
                                    0x8D, 0x52, 0x63, 0x26, 0x52, 0xBF, 0x8D, 0xA0, 0xF0, 0xC1, 0x46, 0xB0, 0x9A, 0x91, 0x17, 0x7E, 
                                    0xCE, 0x2B, 0xDC, 0x93, 0x76, 0x0E, 0x6C, 0x07, 0x42, 0x19, 0x9D, 0x15, 0x54, 0x6E, 0xA7, 0xEE, 
                                    0xD3, 0x69, 0xA5, 0x5B, 0x37, 0x6D, 0x80, 0x4C, 0x9D, 0x18, 0xDD, 0x10, 0xCA, 0x10, 0x36, 0x74, 
                                    0xF0, 0xA1, 0x4A, 0xF7, 0x2E, 0xDF, 0x3A, 0x70, 0x8F, 0x99, 0x60, 0x04, 0x06, 0x9D, 0x3C, 0x75, 
                                    0x1A, 0x5D, 0xF7, 0xB4, 0x0C, 0x60, 0x13, 0x0B, 0xA2, 0x42, 0x47, 0x8F, 0x72, 0x91, 0x56, 0x46, 
                                    0x3E, 0xA3, 0xDD, 0xBF, 0x76, 0x08, 0x09, 0xF2, 0x16, 0x1E, 0xA0, 0xF3, 0xAC, 0x76, 0x1A, 0xC2, 
                                    0x4B, 0x10, 0x80, 0x55, 0x8F, 0x71, 0x0B, 0xC3, 0x62, 0x73, 0xB0, 0xD8, 0xCC, 0xD7, 0x6A, 0x89, 
                                    0xC4, 0x1F, 0x1A, 0xF0, 0x25, 0x6C, 0x6C, 0x9D, 0x7D, 0xC4, 0x30, 0xC5, 0x42, 0x42, 0x24, 0x3E, 
                                    0xD1, 0xE5, 0x15, 0x8A, 0xA6, 0x7D, 0x32, 0xB1, 0x3C, 0x53, 0x78, 0x2B, 0xAA, 0x87, 0x13, 0x51, 
                                    0x2E, 0x40, 0x59, 0x5D, 0x7B, 0xB4, 0x41, 0xB7, 0x72, 0x2A, 0xA7, 0x88, 0xD8, 0xC2, 0xBC, 0x01, 
                                    0x72, 0x41, 0x0D, 0x74, 0x5D, 0x95, 0x67, 0x9D, 0x70, 0x37, 0xDE, 0x9E, 0xCA, 0x04, 0x6C, 0x36, 
                                    0xA6, 0x5D, 0xEF, 0x0F, 0xFC, 0x1A, 0xE0, 0x50, 0x6C, 0x67, 0xE0, 0x24, 0x2E, 0x2A, 0x00, 0xA1, 
                                    0xC6, 0xE8, 0x16, 0x7D, 0x69, 0x54, 0xF4, 0xFD, 0xF1, 0x0C, 0x6F, 0x5C, 0x74, 0xDE, 0xE2, 0x4D, 
                                    0x1C, 0x0F, 0xD5, 0xE9, 0x21, 0x4F, 0xA3, 0xF5, 0x48, 0xEF, 0x54, 0xE5, 0x50, 0xB9, 0xB6, 0x4C, 
                                    0xC1, 0x64, 0xA4, 0x10, 0x2E, 0x48, 0xA4, 0xAF, 0x23, 0x83, 0x72, 0x86, 0xBD, 0xB7, 0x22, 0x3C, 
                                    0x12, 0xE1, 0xBD, 0xA4, 0xF1, 0xEF, 0xAA, 0x66, 0xA8, 0xB1, 0x46, 0xF0, 0x4A, 0x26, 0xAD, 0x47, 
                                    0xE2, 0x4C, 0xF1, 0xEC, 0x3C, 0x18, 0xAD, 0x6E, 0x8E, 0xB8, 0x45, 0xEA, 0xD1, 0x34, 0xAF, 0xE2, 
                                    0xD5, 0xAB, 0xCA, 0x17, 0xC2, 0x45, 0x04, 0x02, 0x76, 0x28, 0xC5, 0x57, 0x05, 0xC6, 0x46, 0x38, 
                                    0x45, 0xD2, 0x00, 0xE3, 0xB1, 0xD9, 0x28, 0x6B, 0x92, 0x7D, 0xEC, 0xA0, 0xB6, 0x56, 0xFA, 0x92, 
                                    0x88, 0xB5, 0x96, 0x2B, 0xB3, 0x60, 0xF4, 0x2A, 0x83, 0x67, 0x6C, 0x21, 0x5B, 0xC8, 0x61, 0x5B, 
                                    0x96, 0x7E, 0x12, 0x39, 0x06, 0xC8, 0xB5, 0xD3, 0x0F, 0xC2, 0x25, 0xB9, 0xD1, 0x39, 0x67, 0x5A, 
                                    0x69, 0x9D, 0xD6, 0x07, 0x0B, 0xBF, 0x66, 0xD6, 0xAF, 0xD1, 0x94, 0x42, 0x2F, 0x47, 0x9F, 0x04, 
                                    0xCA, 0x85, 0xC7, 0x28, 0x36, 0xC1, 0x86, 0x48, 0xF4, 0x05, 0xD9, 0xBE, 0x44, 0x3F, 0x30, 0x8C, 
                                    0xFA, 0x2E, 0xB6, 0x4E, 0x65, 0x56, 0x15, 0x58, 0x56, 0x62, 0x19, 0x5D, 0xD7, 0xFD, 0xE9, 0x66, 
                                    0x59, 0x0E, 0x4B, 0xDD, 0xE9, 0x39, 0x3D, 0x10, 0x55, 0xE1, 0xA5, 0xB9, 0x0B, 0x42, 0x0B, 0x39, 
                                    0x8D, 0x52, 0x63, 0x26, 0x52, 0xBF, 0x8D, 0xA0, 0xF0, 0xC1, 0x46, 0xB0, 0x9A, 0x91, 0x17, 0x7E, 
                                    0xCE, 0x2B, 0xDC, 0x93, 0x76, 0x0E, 0x6C, 0x07, 0x42, 0x19, 0x9D, 0x15, 0x54, 0x6E, 0xA7, 0xEE, 
                                    0xD3, 0x69, 0xA5, 0x5B, 0x37, 0x6D, 0x80, 0x4C, 0x9D, 0x18, 0xDD, 0x10, 0xCA, 0x10, 0x36, 0x74, 
                                    0xF0, 0xA1, 0x4A, 0xF7, 0x2E, 0xDF, 0x3A, 0x70, 0x8F, 0x99, 0x60, 0x04, 0x06, 0x9D, 0x3C, 0x75, 
                                    0x1A, 0x5D, 0xF7, 0xB4, 0x0C, 0x60, 0x13, 0x0B, 0xA2, 0x42, 0x47, 0x8F, 0x72, 0x91, 0x56, 0x46, 
                                    0x3E, 0xA3, 0xDD, 0xBF, 0x76, 0x08, 0x09, 0xF2, 0x16, 0x1E, 0xA0, 0xF3, 0xAC, 0x76, 0x1A, 0xC2, 
                                    0x4B, 0x10, 0x80, 0x55, 0x8F, 0x71, 0x0B, 0xC3, 0x62, 0x73, 0xB0, 0xD8, 0xCC, 0xD7, 0x6A, 0x89, 
                                    0xC4, 0x1F, 0x1A, 0xF0, 0x25, 0x6C, 0x6C, 0x9D, 0x7D, 0xC4, 0x30, 0xC5, 0x42, 0x42, 0x24, 0x3E, 
                                    0xD1, 0xE5, 0x15, 0x8A, 0xA6, 0x7D, 0x32, 0xB1, 0x3C, 0x53, 0x78, 0x2B, 0xAA, 0x87, 0x13, 0x51, 
                                    0x2E, 0x40, 0x59, 0x5D, 0x7B, 0xB4, 0x41, 0xB7, 0x72, 0x2A, 0xA7, 0x88, 0xD8, 0xC2, 0xBC, 0x01, 
                                    0x72, 0x41, 0x0D, 0x74, 0x5D, 0x95, 0x67, 0x9D, 0x70, 0x37, 0xDE, 0x9E, 0xCA, 0x04, 0x6C, 0x36, 
                                    0xA6, 0x5D, 0xEF, 0x0F, 0xFC, 0x1A, 0xE0, 0x50, 0x6C, 0x67, 0xE0, 0x24, 0x2E, 0x2A, 0x00, 0xA1, 
                                    0xC6, 0xE8, 0x16, 0x7D, 0x69, 0x54, 0xF4, 0xFD, 0xF1, 0x0C, 0x6F, 0x5C, 0x74, 0xDE, 0xE2, 0x4D, 
                                    0x1C, 0x0F, 0xD5, 0xE9, 0x21, 0x4F, 0xA3, 0xF5, 0x48, 0xEF, 0x54, 0xE5, 0x50, 0xB9, 0xB6, 0x4C, 
                                    0xC1, 0x64, 0xA4, 0x10, 0x2E, 0x48, 0xA4, 0xAF, 0x23, 0x83, 0x72, 0x86, 0xBD, 0xB7, 0x22, 0x3C, 
                                    0x12, 0xE1, 0xBD, 0xA4, 0xF1, 0xEF, 0xAA, 0x66, 0xA8, 0xB1, 0x46, 0xF0, 0x4A, 0x26, 0xAD, 0x47,
                                    0xE2, 0x4C, 0xF1, 0xEC, 0x3C, 0x18, 0xAD, 0x6E, 0x8E, 0xB8, 0x45, 0xEA, 0xD1, 0x34, 0xAF, 0xE2, 
                                    0xD5, 0xAB, 0xCA, 0x17, 0xC2, 0x45, 0x04, 0x02, 0x76, 0x28, 0xC5, 0x57, 0x05, 0xC6, 0x46, 0x38, 
                                    0x45, 0xD2, 0x00, 0xE3, 0xB1, 0xD9, 0x28, 0x6B, 0x92, 0x7D, 0xEC, 0xA0, 0xB6, 0x56, 0xFA, 0x92, 
                                    0x88, 0xB5, 0x96, 0x2B, 0xB3, 0x60, 0xF4, 0x2A, 0x83, 0x67, 0x6C, 0x21, 0x5B, 0xC8, 0x61, 0x5B, 
                                    0x96, 0x7E, 0x12, 0x39, 0x06, 0xC8, 0xB5, 0xD3, 0x0F, 0xC2, 0x25, 0xB9, 0xD1, 0x39, 0x67, 0x5A, 
                                    0x69, 0x9D, 0xD6, 0x07, 0x0B, 0xBF, 0x66, 0xD6, 0xAF, 0xD1, 0x94, 0x42, 0x2F, 0x47, 0x9F, 0x04, 
                                    0xCA, 0x85, 0xC7, 0x28, 0x36, 0xC1, 0x86, 0x48, 0xF4, 0x05, 0xD9, 0xBE, 0x44, 0x3F, 0x30, 0x8C, 
                                    0xFA, 0x2E, 0xB6, 0x4E, 0x65, 0x56, 0x15, 0x58, 0x56, 0x62, 0x19, 0x5D, 0xD7, 0xFD, 0xE9, 0x66, 
                                    0x59, 0x0E, 0x4B, 0xDD, 0xE9, 0x39, 0x3D, 0x10, 0x55, 0xE1, 0xA5, 0xB9, 0x0B, 0x42, 0x0B, 0x39, 
                                    0x8D, 0x52, 0x63, 0x26, 0x52, 0xBF, 0x8D, 0xA0, 0xF0, 0xC1, 0x46, 0xB0, 0x9A, 0x91, 0x17, 0x7E, 
                                    0xCE, 0x2B, 0xDC, 0x93, 0x76, 0x0E, 0x6C, 0x07, 0x42, 0x19, 0x9D, 0x15, 0x54, 0x6E, 0xA7, 0xEE, 
                                    0xD3, 0x69, 0xA5, 0x5B, 0x37, 0x6D, 0x80, 0x4C, 0x9D, 0x18, 0xDD, 0x10, 0xCA, 0x10, 0x36, 0x74, 
                                    0xF0, 0xA1, 0x4A, 0xF7, 0x2E, 0xDF, 0x3A, 0x70, 0x8F, 0x99, 0x60, 0x04, 0x06, 0x9D, 0x3C, 0x75, 
                                    0x1A, 0x5D, 0xF7, 0xB4, 0x0C, 0x60, 0x13, 0x0B, 0xA2, 0x42, 0x47, 0x8F, 0x72, 0x91, 0x56, 0x46, 
                                    0x3E, 0xA3, 0xDD, 0xBF, 0x76, 0x08, 0x09, 0xF2, 0x16, 0x1E, 0xA0, 0xF3, 0xAC, 0x76, 0x1A, 0xC2, 
                                    0x4B, 0x10, 0x80, 0x55, 0x8F, 0x71, 0x0B, 0xC3, 0x62, 0x73, 0xB0, 0xD8, 0xCC, 0xD7, 0x6A, 0x89, 
                                    0xC4, 0x1F, 0x1A, 0xF0, 0x25, 0x6C, 0x6C, 0x9D, 0x7D, 0xC4, 0x30, 0xC5, 0x42, 0x42, 0x24, 0x3E, 
                                    0xD1, 0xE5, 0x15, 0x8A, 0xA6, 0x7D, 0x32, 0xB1, 0x3C, 0x53, 0x78, 0x2B, 0xAA, 0x87, 0x13, 0x51, 
                                    0x2E, 0x40, 0x59, 0x5D, 0x7B, 0xB4, 0x41, 0xB7, 0x72, 0x2A, 0xA7, 0x88, 0xD8, 0xC2, 0xBC, 0x01, 
                                    0x72, 0x41, 0x0D, 0x74, 0x5D, 0x95, 0x67, 0x9D, 0x70, 0x37, 0xDE, 0x9E, 0xCA, 0x04, 0x6C, 0x36, 
                                    0xA6, 0x5D, 0xEF, 0x0F, 0xFC, 0x1A, 0xE0, 0x50, 0x6C, 0x67, 0xE0, 0x24, 0x2E, 0x2A, 0x00, 0xA1, 
                                    0xC6, 0xE8, 0x16, 0x7D, 0x69, 0x54, 0xF4, 0xFD, 0xF1, 0x0C, 0x6F, 0x5C, 0x74, 0xDE, 0xE2, 0x4D, 
                                    0x1C, 0x0F, 0xD5, 0xE9, 0x21, 0x4F, 0xA3, 0xF5, 0x48, 0xEF, 0x54, 0xE5, 0x50, 0xB9, 0xB6, 0x4C, 
                                    0xC1, 0x64, 0xA4, 0x10, 0x2E, 0x48, 0xA4, 0xAF, 0x23, 0x83, 0x72, 0x86, 0xBD, 0xB7, 0x22, 0x3C, 
                                    0x12, 0xE1, 0xBD, 0xA4, 0xF1, 0xEF, 0xAA, 0x66, 0xA8, 0xB1, 0x46, 0xF0, 0x4A, 0x26, 0xAD, 0x47, 
                                    0xE2, 0x4C, 0xF1, 0xEC, 0x3C, 0x18, 0xAD, 0x6E, 0x8E, 0xB8, 0x45, 0xEA, 0xD1, 0x34, 0xAF, 0xE2, 
                                    0xD5, 0xAB, 0xCA, 0x17, 0xC2, 0x45, 0x04, 0x02, 0x76, 0x28, 0xC5, 0x57, 0x05, 0xC6, 0x46, 0x38, 
                                    0x45, 0xD2, 0x00, 0xE3, 0xB1, 0xD9, 0x28, 0x6B, 0x92, 0x7D, 0xEC, 0xA0, 0xB6, 0x56, 0xFA, 0x92, 
                                    0x88, 0xB5, 0x96, 0x2B, 0xB3, 0x60, 0xF4, 0x2A, 0x83, 0x67, 0x6C, 0x21, 0x5B, 0xC8, 0x61, 0x5B, 
                                    0x96, 0x7E, 0x12, 0x39, 0x06, 0xC8, 0xB5, 0xD3, 0x0F, 0xC2, 0x25, 0xB9, 0xD1, 0x39, 0x67, 0x5A, 
                                    0x69, 0x9D, 0xD6, 0x07, 0x0B, 0xBF, 0x66, 0xD6, 0xAF, 0xD1, 0x94, 0x42, 0x2F, 0x47, 0x9F, 0x04, 
                                    0xCA, 0x85, 0xC7, 0x28, 0x36, 0xC1, 0x86, 0x48, 0xF4, 0x05, 0xD9, 0xBE, 0x44, 0x3F, 0x30, 0x8C, 
                                    0xFA, 0x2E, 0xB6, 0x4E, 0x65, 0x56, 0x15, 0x58, 0x56, 0x62, 0x19, 0x5D, 0xD7, 0xFD, 0xE9, 0x66, 
                                    0x59, 0x0E, 0x4B, 0xDD, 0xE9, 0x39, 0x3D, 0x10, 0x55, 0xE1, 0xA5, 0xB9, 0x0B, 0x42, 0x0B, 0x39, 
                                    0x8D, 0x52, 0x63, 0x26, 0x52, 0xBF, 0x8D, 0xA0, 0xF0, 0xC1, 0x46, 0xB0, 0x9A, 0x91, 0x17, 0x7E, 
                                    0xCE, 0x2B, 0xDC, 0x93, 0x76, 0x0E, 0x6C, 0x07, 0x42, 0x19, 0x9D, 0x15, 0x54, 0x6E, 0xA7, 0xEE, 
                                    0xD3, 0x69, 0xA5, 0x5B, 0x37, 0x6D, 0x80, 0x4C, 0x9D, 0x18, 0xDD, 0x10, 0xCA, 0x10, 0x36, 0x74, 
                                    0xF0, 0xA1, 0x4A, 0xF7, 0x2E, 0xDF, 0x3A, 0x70, 0x8F, 0x99, 0x60, 0x04, 0x06, 0x9D, 0x3C, 0x75, 
                                    0x1A, 0x5D, 0xF7, 0xB4, 0x0C, 0x60, 0x13, 0x0B, 0xA2, 0x42, 0x47, 0x8F, 0x72, 0x91, 0x56, 0x46, 
                                    0x3E, 0xA3, 0xDD, 0xBF, 0x76, 0x08, 0x09, 0xF2, 0x16, 0x1E, 0xA0, 0xF3, 0xAC, 0x76, 0x1A, 0xC2, 
                                    0x4B, 0x10, 0x80, 0x55, 0x8F, 0x71, 0x0B, 0xC3, 0x62, 0x73, 0xB0, 0xD8, 0xCC, 0xD7, 0x6A, 0x89, 
                                    0xC4, 0x1F, 0x1A, 0xF0, 0x25, 0x6C, 0x6C, 0x9D, 0x7D, 0xC4, 0x30, 0xC5, 0x42, 0x42, 0x24, 0x3E, 
                                    0xD1, 0xE5, 0x15, 0x8A, 0xA6, 0x7D, 0x32, 0xB1, 0x3C, 0x53, 0x78, 0x2B, 0xAA, 0x87, 0x13, 0x51, 
                                    0x2E, 0x40, 0x59, 0x5D, 0x7B, 0xB4, 0x41, 0xB7, 0x72, 0x2A, 0xA7, 0x88, 0xD8, 0xC2, 0xBC, 0x01, 
                                    0x72, 0x41, 0x0D, 0x74, 0x5D, 0x95, 0x67, 0x9D, 0x70, 0x37, 0xDE, 0x9E, 0xCA, 0x04, 0x6C, 0x36, 
                                    0xA6, 0x5D, 0xEF, 0x0F, 0xFC, 0x1A, 0xE0, 0x50, 0x6C, 0x67, 0xE0, 0x24, 0x2E, 0x2A, 0x00, 0xA1, 
                                    0xC6, 0xE8, 0x16, 0x7D, 0x69, 0x54, 0xF4, 0xFD, 0xF1, 0x0C, 0x6F, 0x5C, 0x74, 0xDE, 0xE2, 0x4D, 
                                    0x1C, 0x0F, 0xD5, 0xE9, 0x21, 0x4F, 0xA3, 0xF5, 0x48, 0xEF, 0x54, 0xE5, 0x50, 0xB9, 0xB6, 0x4C, 
                                    0xC1, 0x64, 0xA4, 0x10, 0x2E, 0x48, 0xA4, 0xAF, 0x23, 0x83, 0x72, 0x86, 0xBD, 0xB7, 0x22, 0x3C, 
                                    0x12, 0xE1, 0xBD, 0xA4, 0xF1, 0xEF, 0xAA, 0x66, 0xA8, 0xB1, 0x46, 0xF0, 0x4A, 0x26, 0xAD, 0x47};
uint8_t exampleKey[KEY_SIZE] = {'A', '0', 'B', '1', 'C', '2', 'D', '3', 
                                'E', '4', 'F', '5', 'G', '6', 'H', '7', 
                                'I', '8', 'J', '9', 'K', '0', 'L', '1', 
                                'M', '2', 'N', '3', 'O', '4', 'P', '5',};

MemRef<unsigned char, 1> input{exampleInput, INPUT_SIZE};
MemRef<unsigned char, 1> key{exampleKey, KEY_SIZE};
uint8_t *exampleOutput = new uint8_t[OUTPUT_SIZE];
MemRef<unsigned char, 1> output{exampleOutput, OUTPUT_SIZE};

void printResult(uint8_t* output) {
    // Print the output.
    std::cout << "[";
    for (size_t i = 0; i < AES_DIGEST_LEN; i++) {
      printf("%c", output[i]);
    }
    std::cout << "]" << std::endl;
}

void printInput(int idx, uint8_t* input) {
  // Print the input.
  std::cout << "Input " << idx << ": [0x";
  for (size_t i = 0; i < MSG_LEN; i++) {
    printf("%02X", input[i]);
  }
  std::cout << "] ==> ";
}

void generateResultGPUAES() {
  int keyLen = static_cast<int>(AES_TYPE::AES128);
  
#ifdef WAFER_OPENSSL
  std::cout << "--------------------------------------------------------\nOpenSSL AES\n";

  const clock_t begin_time = clock();

for (int i = 0; i < MSG_CNT; i++) {
  printInput(i, exampleInput + (i * MSG_LEN));

  EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();

  const EVP_CIPHER *cipher;
  if (keyLen == 16)
    cipher = EVP_aes_128_ecb();
  else if (keyLen == 24)
    cipher = EVP_aes_192_ecb();
  else // keyLen == 32
    cipher = EVP_aes_256_ecb();

  EVP_DecryptInit_ex(ctx, cipher, NULL, exampleKey, NULL);
  int outlen;
  int plaintext_len = 0;
  EVP_DecryptUpdate(ctx, exampleOutput + (i * AES_DIGEST_LEN), &outlen, exampleInput + (i * MSG_LEN), AES_DIGEST_LEN);
  plaintext_len += outlen;
  EVP_DecryptFinal_ex(ctx, exampleOutput + (i * AES_DIGEST_LEN) + plaintext_len, &outlen);
  plaintext_len += outlen;

  printResult(exampleOutput + (i * AES_DIGEST_LEN));
  EVP_CIPHER_CTX_free(ctx);
}

  float duration = float(clock() - begin_time) / CLOCKS_PER_SEC;
  std::cout << "Calculate duration: " << std::fixed << duration << std::endl;
  std::cout << "--------------------------------------------------------"
            << std::endl;
#else
  const clock_t begin_time = clock();
  #pragma wafer parallel
  AES_CRACK_OP(&input, MSG_LEN, &output, AES_DIGEST_LEN, &key, keyLen, MSG_CNT);
  float duration = float(clock() - begin_time) / CLOCKS_PER_SEC;

  std::cout << "\n--------------------------------------------------------\nMLIR AES\n";
  
  for (int i = 0; i < MSG_CNT; i++) {
    printInput(i, exampleInput + (i * MSG_LEN));
    printResult(exampleOutput + (i * AES_DIGEST_LEN));
  }

  std::cout << "Key " << ": [";
  for (size_t i = 0; i < keyLen; i++) {
    std::cout << key[i];
  }
  std::cout << "]" << std::endl;
  std::cout << "Calculate duration: " << std::fixed << duration << std::endl;
  std::cout << "--------------------------------------------------------"
            << std::endl;
#endif
}

int main() {
  generateResultGPUAES();
  return 0;
}
